# -*- coding: utf-8 -*-
# @Time    : 2021/10/18 11:19
# @Author  : huangwei
# @File    : get_cell2.py
# @Software: PyCharm
import math

import cv2
import numpy as np
from skimage import measure

from method import draw_box
from table_method import get_line, draw_lines, sqrt, line_line, roll_min_rect_box, minAreaRectBox, dist
from table_net import table_net
import tensorflow as tf

config = tf.compat.v1.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
config.gpu_options.allow_growth = True
session = tf.compat.v1.InteractiveSession(config=config)

tableModeLinePath = 'models/table-line.h5'
model = table_net((None, None, 3), 2)
model.load_weights(tableModeLinePath)


def resize( img, size ):
    img_h, img_w = img.shape[:2]
    w, h = size

    ratio = min(w / img_w, h / img_h)
    new_w = int(img_w * ratio)
    new_h = int(img_h * ratio)

    resize_img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
    cv2.imwrite('temp_path/resize_img.png', resize_img)

    fill_value = [128, 128, 128]
    boxed_image = np.zeros((h, w, 3), dtype=np.uint8)
    boxed_image[:] = fill_value
    boxed_image[:new_h, :new_w, :] = resize_img

    # 分别为长、宽变化的比例
    fx = new_w / img_w
    fy = new_h / img_h

    return boxed_image, fx, fy


def adjust_lines( RowsLines, ColsLines, alph=50 ):
    ##调整line，合并相交线
    nrow = len(RowsLines)
    ncol = len(ColsLines)
    newRowsLines = []
    newColsLines = []
    for i in range(nrow):

        x1, y1, x2, y2 = RowsLines[i]
        cx1, cy1 = (x1 + x2) / 2, (y1 + y2) / 2
        for j in range(nrow):
            if i != j:
                x3, y3, x4, y4 = RowsLines[j]
                cx2, cy2 = (x3 + x4) / 2, (y3 + y4) / 2
                if (x3 < cx1 < x4 or y3 < cy1 < y4) or (x1 < cx2 < x2 or y1 < cy2 < y2):
                    continue
                else:
                    r = sqrt((x1, y1), (x3, y3))
                    if r < alph:
                        newRowsLines.append([x1, y1, x3, y3])
                    r = sqrt((x1, y1), (x4, y4))
                    if r < alph:
                        newRowsLines.append([x1, y1, x4, y4])

                    r = sqrt((x2, y2), (x3, y3))
                    if r < alph:
                        newRowsLines.append([x2, y2, x3, y3])
                    r = sqrt((x2, y2), (x4, y4))
                    if r < alph:
                        newRowsLines.append([x2, y2, x4, y4])

    for i in range(ncol):
        x1, y1, x2, y2 = ColsLines[i]
        cx1, cy1 = (x1 + x2) / 2, (y1 + y2) / 2
        for j in range(ncol):
            if i != j:
                x3, y3, x4, y4 = ColsLines[j]
                cx2, cy2 = (x3 + x4) / 2, (y3 + y4) / 2
                if (x3 < cx1 < x4 or y3 < cy1 < y4) or (x1 < cx2 < x2 or y1 < cy2 < y2):
                    continue
                else:
                    r = sqrt((x1, y1), (x3, y3))
                    if r < alph:
                        newColsLines.append([x1, y1, x3, y3])
                    r = sqrt((x1, y1), (x4, y4))
                    if r < alph:
                        newColsLines.append([x1, y1, x4, y4])

                    r = sqrt((x2, y2), (x3, y3))
                    if r < alph:
                        newColsLines.append([x2, y2, x3, y3])
                    r = sqrt((x2, y2), (x4, y4))
                    if r < alph:
                        newColsLines.append([x2, y2, x4, y4])

    return newRowsLines, newColsLines


def get_table_line( pred_line, axis, alph ):
    # axis=0 横线
    # axis=1 竖线
    lables = measure.label(pred_line > 0, connectivity=2)
    regions = measure.regionprops(lables)

    if axis == 1:
        lines = [get_line(line.coords) for line in regions if line.bbox[2] - line.bbox[0] > alph]
    else:
        lines = [get_line(line.coords) for line in regions if line.bbox[3] - line.bbox[1] > alph]

    return lines


# 去除不平行的线
def not_par_line( lines, alph=1 ):
    degree_list = []
    line_num = len(lines)
    for line in lines:
        x1, y1, x2, y2 = line
        if x1 != x2:
            k12 = (y2 - y1) / (x2 - x1)
            a12 = math.degrees(math.atan(k12))
        else:
            a12 = 90
        degree_list.append(a12)

    sort_list = sorted(degree_list)
    mid_degree = sort_list[line_num // 2]

    new_lines = []
    for i in range(line_num):
        if abs(degree_list[i] - mid_degree) < alph:
            new_lines.append(lines[i])

    return new_lines


def find_unclear_lines( lines, alph=15 ):
    # 找出不清楚的线并且合并线段
    line_num = len(lines)

    new_lines = []
    line_used = []

    # 四个点在一条直线上且最近的两个点距离小于 alph
    if line_num > 1:
        for i in range(line_num - 1):
            x1, y1, x2, y2 = lines[i]
            if x1 != x2:
                k12 = (y2 - y1) / (x2 - x1)
                a12 = math.degrees(math.atan(k12))
            else:
                a12 = 90

            for j in range(i + 1, line_num):
                x3, y3, x4, y4 = lines[j]
                if x3 != x4:
                    k34 = (y4 - y3) / (x4 - x3)
                    a34 = math.degrees(math.atan(k34))
                else:
                    a34 = 90

                # print("a12:{}, a34:{}".format(a12, a34))
                # 如果两条线的角度偏差不超过2°则认为是平行的
                if abs(a34 - a12) < 2:
                    # 计算 a13
                    if x1 != x3:
                        k13 = (y3 - y1) / (x3 - x1)
                        a13 = math.degrees(math.atan(k13))
                    else:
                        a13 = 90

                    # 如果 a12 和 a13相差不超过1°则认为 12和13重合
                    # 求四个点之间的距离，最短距离小于alph则连接，用最长的那条取代
                    if abs(a13 - a12) < 1:
                        dis13 = dist((x1, y1), (x3, y3))
                        dis14 = dist((x1, y1), (x4, y4))
                        dis23 = dist((x2, y2), (x3, y3))
                        dis24 = dist((x2, y2), (x4, y4))
                        dist_list = [dis13, dis14, dis23, dis24]
                        min_dis = min(dist_list)
                        # print("min dis is:", min_dis)

                        if min_dis < alph:
                            max_index = dist_list.index(max(dist_list))
                            # print("max index:", max_index)
                            if max_index == 0:
                                # 返回13
                                new_lines.append([x1, y1, x3, y3])
                            elif max_index == 1:
                                # 返回14
                                new_lines.append([x1, y1, x4, y4])
                            elif max_index == 2:
                                # 返回23
                                new_lines.append([x2, y2, x3, y3])
                            else:
                                # 返回24
                                new_lines.append([x2, y2, x4, y4])

                            # 记录已经使用的线段
                            # print("use:", (i, j))
                            line_used.append(i)
                            line_used.append(j)

    for i in range(line_num):
        if i not in line_used:
            new_lines.append(lines[i])

    return new_lines


def sort_line( lines, axis=0 ):
    center_list = []
    for line in lines:
        x1, y1, x2, y2 = line
        c_x, c_y = (x1 + x2) / 2, (y1 + y2) / 2

        if axis == 0:
            center_list.append(c_y)
        else:
            center_list.append(c_x)

    sort_list = np.argsort(center_list)

    new_lines = []
    for index in sort_list:
        new_lines.append(lines[index])

    return new_lines


def is_intersect( line1, line2 ):
    x1, y1, x2, y2 = line1
    x3, y3, x4, y4 = line2

    # 3，4是否在12两侧
    x12 = x2 - x1
    y12 = y2 - y1
    x13 = x3 - x1
    y13 = y3 - y1
    x14 = x4 - x1
    y14 = y4 - y1
    d1 = x12 * y13 - x13 * y12
    d2 = x12 * y14 - x14 * y12
    delta1 = d1 * d2

    # 1，2是否在34两侧
    x34 = x4 - x3
    y34 = y4 - y3
    x31 = x1 - x3
    y31 = y1 - y3
    x32 = x2 - x3
    y32 = y2 - y3
    d3 = x34 * y31 - x31 * y34
    d4 = x34 * y32 - x32 * y34
    delta2 = d3 * d4

    if delta1 <= 0 and delta2 <= 0:
        return True
    else:
        return False


def get_intersect_index( line, lines_list ):
    line_num = len(lines_list)
    cross_list = []

    for j in range(line_num):
        # 判断两条线段是否相交
        is_inter = is_intersect(line, lines_list[j])
        if is_inter:
            cross_list.append(j)

    return cross_list


def get_intersect_point( line1, line2 ):
    # 求两条线段的交点
    x1, y1, x2, y2 = line1
    x3, y3, x4, y4 = line2

    px = ((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / (
            (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4))
    py = ((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / (
            (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4))

    return px, py


def fit_lines( lines1, lines2, axis=0, is_add=True ):
    new_lines = []
    for line1 in lines1:
        cross_list = get_intersect_index(line1, lines2)
        # 如果相交条数小于2，则删除这条线
        if len(cross_list) > 1:
            # 求第一个和最后一个相交的交点
            temp_line1 = lines2[cross_list[0]]
            temp_line2 = lines2[cross_list[-1]]

            # 求两条线段的交点
            cx1, cy1 = get_intersect_point(line1, temp_line1)
            cx2, cy2 = get_intersect_point(line1, temp_line2)
            x1, y1, x2, y2 = line1
            dis1 = dist((x1, y1), (cx1, cy1))
            dis2 = dist((x2, y2), (cx2, cy2))

            # 如果突出部分小于五个像素或为最末端则删除这部分，否则
            if dis1 < 5 or cross_list[0] == 0:
                if is_add:
                    if axis == 0:
                        x1, y1 = cx1 - 1, cy1
                    else:
                        x1, y1 = cx1, cy1 - 1
                else:
                    x1, y1 = cx1, cy1
            else:
                # 需要添加突出部分过长应该怎么解决
                pass

            if dis2 < 5 or cross_list[-1] == (len(lines2) - 1):
                if is_add:
                    if axis == 0:
                        x2, y2 = cx2 + 1, cy2
                    else:
                        x2, y2 = cx2, cy2 + 1
                else:
                    x2, y2 = cx2, cy2
            else:
                pass

            line1 = [x1, y1, x2, y2]
            new_lines.append(line1)
        else:
            # 直接删除这条线不添加到新的list中
            pass

    return new_lines


def get_lines( img, input_size, prob=0.2, alph=15, row=30, col=30 ):
    resize_img, fx, fy = resize(img, input_size)

    pred = model.predict(np.array([np.array(resize_img) / 255.0]))
    pred = pred[0]

    # test_row = pred[..., 0]
    # test_col = pred[..., 1]

    rows = pred[..., 0] > prob
    cols = pred[..., 1] > prob
    rows = rows.astype(int)
    cols = cols.astype(int)

    # 根据模型得出的结果提取出识别出的横线和竖线
    row_lines = get_table_line(rows, axis=0, alph=row)
    col_lines = get_table_line(cols, axis=1, alph=col)

    # 将提取出的线画在黑板上
    tmp = np.zeros(input_size, dtype='uint8')
    tmp = draw_lines(tmp, row_lines + col_lines, color=255, line_width=2)
    cv2.imwrite('temp_path/aa.png', tmp)

    # 去除识别出来的线中不平行的线
    row_lines = not_par_line(row_lines)
    col_lines = not_par_line(col_lines)

    # 断续的横线或竖线进行连接
    row_lines = find_unclear_lines(row_lines, alph)
    col_lines = find_unclear_lines(col_lines, alph)

    tmp = np.zeros(input_size, dtype='uint8')
    tmp = draw_lines(tmp, row_lines + col_lines, color=255, line_width=2)
    cv2.imwrite('temp_path/bb.png', tmp)

    # 转换成原尺寸
    if len(row_lines) > 0:
        row_lines = np.array(row_lines)
        row_lines[:, [0, 2]] = row_lines[:, [0, 2]] / fx
        row_lines[:, [1, 3]] = row_lines[:, [1, 3]] / fy
        row_lines = row_lines.tolist()

    if len(col_lines) > 0:
        col_lines = np.array(col_lines)
        col_lines[:, [0, 2]] = col_lines[:, [0, 2]] / fx
        col_lines[:, [1, 3]] = col_lines[:, [1, 3]] / fy
        col_lines = col_lines.tolist()

    # 先对线进行从左到右从上到下排序
    row_lines = sort_line(row_lines)
    col_lines = sort_line(col_lines, axis=1)

    # 依次扫描每一条横线和竖线，去除冒出头的部分
    row_lines = fit_lines(row_lines, col_lines)
    col_lines = fit_lines(col_lines, row_lines, axis=1)
    row_lines = fit_lines(row_lines, col_lines)
    # row_lines = fit_lines(row_lines, col_lines, is_add=False)
    # col_lines = fit_lines(col_lines, row_lines, axis=1, is_add=False)

    return row_lines, col_lines

# file_path = "images/simple3.png"
# img = cv2.imread(file_path)
# input_size = (1024, 1024)
# rows_lines, cols_lines = get_lines(img, input_size)
#
# # 找出格子
# # 将找到的所有线都在全黑的图上画出，再在这个图上寻找连通区域即每一个格子
# tmp = np.zeros((img.shape[0], img.shape[1]), dtype='uint8')
# tmp = draw_lines(tmp, rows_lines + cols_lines, color=255, line_width=1)
# tmp_filepath = 'temp_path/tmp.png'
# cv2.imwrite(tmp_filepath, tmp)
#
# table_labels = measure.label(tmp == 0, connectivity=2)
# regions = measure.regionprops(table_labels)
# rboxes = []
# for region in regions:
#     # print("area:", region.bbox_area)
#     if region.bbox_area < img.shape[0] * img.shape[1] / 2:
#         rbox = minAreaRectBox(region.coords)
#         rboxes.append(rbox)
#
# for box in rboxes:
#     draw_box(img, box)
# cv2.imwrite("temp_path/box.png", img)
# print(len(rboxes))
